# -- encoding:utf-8 --

from flask import Flask, request, jsonify

import sys
import numpy as np
import tensorflow as tf
from utils.vocabulary_utils import VocabularyProcessorUtil


class TextClassifyParser(object):
    def __init__(self, pb_path, return_elements, vocab_model_path, label_names):
        self.label_names = label_names
        tf.logging.info("开始进行TensorFlow模型恢复操作....")
        with tf.Graph().as_default():
            # 1. 构建图和会话
            self.graph = tf.get_default_graph()
            self.sess = tf.Session(graph=self.graph)

            # 2. 读取pb文件内容（二进制字符串形式）并恢复图
            with tf.gfile.GFile(pb_path, "rb") as reader:
                frozen_graph_def = tf.GraphDef()
                frozen_graph_def.ParseFromString(reader.read())

            # 3. 进行执行图以及模型参数的恢复, 并且得到关心的Tensor对象
            tensors = tf.import_graph_def(graph_def=frozen_graph_def, return_elements=return_elements)
            print(tensors)

            # 4. 获取关心的Tensor对象
            self.inputs = tensors[0]
            self.predictions = tensors[1]
            self.probability = tensors[2]
        tf.logging.info("TensorFlow模型恢复完成!!!!")
        tf.logging.info("开始恢复词向量转换模型.....")
        if not tf.gfile.Exists(vocab_model_path):
            raise Exception("词汇转换模型必须存在,请检查磁盘路径:{}".format(vocab_model_path))
        self.vocab_model = VocabularyProcessorUtil.load_model(save_path=vocab_model_path)
        tf.logging.info("词汇向量转换模型恢复完成!!!!")

    def predict(self, text):
        """
        预测函数
        :param text: 一个数组，eg: ["xxxx", "xxxx"]或者一个字符串
        :return:
        """
        # 0. 判断一下数据格式
        is_one_sample = False
        if isinstance(text, str):
            text = [text]
            is_one_sample = True

        # 1. 将文本转换为单词序列id
        texts = np.asarray(list(self.vocab_model.transform(text)))

        # 2. 调用session的run方法获取预测结果
        predict_id, predict_prob = self.sess.run([self.predictions, self.probability],
                                                 feed_dict={self.inputs: texts})
        # 3. 结果返回
        result = []
        for id, prob in zip(predict_id, predict_prob):
            result.append([id, self.label_names[id], prob[id]])
        return result[0] if is_one_sample else result


if __name__ == '__main__':
    # 获取参数
    port = 8081
    pb_path = "./hotel.model.pb"
    vocab_model_path = "../hotel/model/vocab.pkl"
    label_names = ['差评', '好评']
    if len(sys.argv) > 1:
        args = sys.argv[1:]
        args_size = len(args)
        idx = 0
        while idx < args_size:
            cur_arg = args[idx]
            if cur_arg == '--port':
                port = int(args[idx + 1])
            elif cur_arg == '--pb_path':
                pb_path = args[idx + 1].strip()
            elif cur_arg == '--vocab_path':
                vocab_model_path = args[idx + 1].strip()
            elif cur_arg == '--label_names':
                label_names = []
                for name in args[idx + 1].strip().split(","):
                    label_names.append(name.strip())
            idx += 1
    print("监听端口号为:{}".format(port))
    print("加载的TF模型文件路径为:{}".format(pb_path))
    print("加载的词汇转换模型文件路径为:{}".format(vocab_model_path))
    print("类别数组为:{}".format(label_names))

    # 恢复模型(方式一：恢复ckpt格式的模型)
    network_name = "TEXT_CNN"
    parser = TextClassifyParser(
        pb_path=pb_path,
        return_elements=[
            "{}/placeholders/input_word_id:0".format(network_name),
            "{}/project/predictions:0".format(network_name),
            "{}/project/probability:0".format(network_name)
        ],
        vocab_model_path=vocab_model_path,
        label_names=label_names
    )

    # 一、构建Flask的应用APP
    app = Flask(__name__)
    app.config['JSON_AS_ASCII'] = False  # 返回JSON格式的时候，数据就不会进行编码


    @app.route("/")
    @app.route("/index")
    def index():
        return "欢迎您进入情感分类预测的服务端!!!"


    @app.route("/predict", methods=['GET', 'POST'])
    def predict():
        # 1. 获取参数(待预测的文本数据); 如果没有传入，那么默认为None
        if request.method == 'GET':
            text = request.args.get("text", None)
        else:
            text = request.form.get("text", None)

        # 2. 对text数据进行预测，得到预测结果
        pred = parser.predict(text)

        # 3. 结果数据处理并返回
        result = {
            'code': 200,  # 一般是要给数字用于表示调用返回的结果情况
            'msg': '成功!!!',
            'data': [{
                "class_id": int(pred[0]),
                "class_name": pred[1],
                "probability": float(pred[2])
            }],
            'text': text
        }

        # 4. 以json的格式返回
        return jsonify(result)


    # 二、启动Flask应用
    app.run(host="0.0.0.0", port=port)
